import pandas as pd
import os
import xml.etree.ElementTree as ET

# 设定你的ACE2005文件地址（Chinese）
ace_path = 'D:/课件/NLP/ChineseACE2005/Chinese/'

if __name__ == '__main__':
    with open('D:/课件/NLP/apf.xml-file-in-ACE2005-to-ann/data_list.csv', mode='r', encoding='utf-8') as csv_file:
        rows = csv_file.readlines()
        for row in rows[1:]:
            items = row.replace('\n', '').split(',')
            name = items[0]
            path = os.path.join(ace_path, name)
            sgm_path = (path + '.sgm')

            # print(xml_path)
            # 解析xml文件
            xml_path = (path + '.apf.xml')
            apf_tree = ET.parse(xml_path)
            apf_root = apf_tree.getroot()

            # 获取文件名
            document = apf_root.find('document')
            document_name = document.attrib

            # 创建ann文件
            output = open(path+'.ann', 'w', encoding='utf-8')

            # T的下标，和E的下标。
            # 我们暂时设定实体和论元都是T，触发词是E，便于区分。
            T_index = 1
            E_index = 1

            # 调整位置
            entity1 = [1000, ""]
            for child_document in document:
                # 找到位置最靠前的标记
                entity = child_document
                for child_entity in entity:
                    if child_entity.tag == 'entity_mention':
                        entity_mention = child_entity
                        for head in entity_mention.iter('head'):
                            for entity_charseq in head:  # 遍历 charseq 子节点
                                if entity1[0] > int(entity_charseq.attrib['START']):
                                    entity1[0] = int(entity_charseq.attrib['START'])
                                    entity1[1] = entity_charseq.text

            txt = open(sgm_path, 'r', encoding='utf-8')
            # 找到txt中最靠前的标记，计算需调整距离
            pos = 0
            ifHeadline = 0
            for line in txt:
                if "<HEADLINE>" in line:
                    ifHeadline = 1
                if entity1[1] in line:
                    if ifHeadline:
                        ifHeadline = 0
                    else:
                        pos = pos + line.index(entity1[1])
                        break
                for ch in line:
                    pos += 1
            dis = int(pos - entity1[0])  # 要调整的距离

            for child_document in document:
                # 提取实体
                if child_document.tag == 'entity':
                    entity = child_document
                    entity_type = entity.attrib.get('TYPE')
                    for child_entity in entity:
                        if child_entity.tag == 'entity_mention':
                            entity_mention = child_entity
                            entity_mention_ID = 'T' + str(T_index)
                            T_index += 1
                            for head in entity_mention.iter('head'):
                                for entity_charseq in head:  # 遍历 charseq 子节点
                                    entity_mention_start = str(dis + int(entity_charseq.attrib['START']))
                                    entity_mention_end = str(dis + int(entity_charseq.attrib['END']) + 1)
                                    # 设置跨行格式
                                    pos = 0
                                    if '\n' in entity_charseq.text:
                                        pos = entity_charseq.text.index('\n')
                                        entity_charseq.text = entity_charseq.text.replace('\n', ' ')
                                        entity_mention_end = str(
                                            int(entity_mention_start) + pos + 1) + ' ' + entity_mention_end
                                        entity_mention_start = entity_mention_start + ' ' + str(
                                            int(entity_mention_start) + pos)
                                        output.write(
                                            entity_mention_ID + '\t' + entity_type + ' ' + entity_mention_start + ';' +
                                            entity_mention_end + '\t' + entity_charseq.text+'\n')
                                    else:
                                        output.write(
                                            entity_mention_ID + '\t' + entity_type + ' ' + entity_mention_start + ' ' +
                                            entity_mention_end + '\t' + entity_charseq.text+'\n')

                # 提取事件
                elif child_document.tag == 'event':
                    event = child_document
                    event_type = event.attrib.get('TYPE')
                    event_ID = event.attrib.get('ID')
                    event_argument_role = []
                    event_argument_ID = []
                    event_argument = {}
                    for child_event in event:
                        # 提取 event_argument 属性
                        if child_event.tag == 'event_argument':
                            event_argument_role.append(child_event.get('ROLE'))
                            event_argument_ID.append(child_event.get('REFID'))

                        # 提取 event_mention
                        elif child_event.tag == 'event_mention':
                            event_mention = child_event
                            event_mention_id = event_mention.attrib.get('ID')
                            for child_event_mention in event_mention:
                                # 提取触发词
                                if child_event_mention.tag == 'anchor':
                                    trigger = child_event_mention  # trigger 的 id 是event_mention 的 id
                                    trigger_ID = 'E' + str(E_index)
                                    E_index += 1
                                    for trigger_charseq in trigger:
                                        trigger_start = str(dis + int(trigger_charseq.attrib['START']))
                                        trigger_end = str(dis + int(trigger_charseq.attrib['END']) + 1)
                                        # 设置跨行格式
                                        pos = 0
                                        if '\n' in trigger_charseq.text:
                                            pos = trigger_charseq.text.index('\n')
                                            trigger_charseq.text = trigger_charseq.text.replace('\n', ' ')
                                            trigger_mention_end = str(
                                                int(trigger_start) + pos + 1) + ' ' + trigger_end
                                            trigger_start = trigger_start + ' ' + str(
                                                int(trigger_start) + pos)
                                            output.write(
                                                trigger_ID + '\t' + event_type + ' ' + trigger_start + ';' +
                                                trigger_end + '\t' + trigger_charseq.text+'\n')
                                        else:
                                            output.write(
                                                trigger_ID + '\t' + event_type + ' ' + trigger_start + ' ' +
                                                trigger_end + '\t' + trigger_charseq.text+'\n')
                                        # print(
                                        #   trigger_ID + '\t' + event_type + ' ' + trigger_start + ' ' +
                                        #   trigger_end + '\t' + trigger_charseq.text)

                                # 提取 event_mention 中的 argument
                                elif child_event_mention.tag == 'event_mention_argument':
                                    argument = child_event_mention
                                    argument_ID = 'T' + str(T_index)
                                    T_index += 1
                                    argument_role = argument.attrib['ROLE']

                                    for argument_charseq in argument.iter('charseq'):
                                        argument_start = str(dis + int(argument_charseq.attrib['START']))
                                        argument_end = str(dis + int(argument_charseq.attrib['END']) + 1)
                                        # 设置跨行格式
                                        pos = 0
                                        if '\n' in argument_charseq.text:
                                            pos = argument_charseq.text.index('\n')
                                            argument_charseq.text = argument_charseq.text.replace('\n', ' ')
                                            argument_end = str(
                                                int(argument_start) + pos + 1) + ' ' + argument_end
                                            argument_start = argument_start + ' ' + str(
                                                int(argument_start) + pos)
                                            output.write(
                                                argument_ID + '\t' + argument_role + ' ' + argument_start + ';' +
                                                argument_end + '\t' + argument_charseq.text+'\n')
                                        else:
                                            output.write(
                                                argument_ID + '\t' + argument_role + ' ' + argument_start + ' ' +
                                                argument_end + '\t' + argument_charseq.text+'\n')
                                        #print(
                                        #    argument_ID + '\t' + argument_role + ' ' + argument_start + ' ' +
                                        #    argument_end + '\t' + argument_charseq.text)

                            # print(event_ID + '\t', end='')
                            # for a,b in zip(event_argument_ID,event_argument_role):
                            # c=str(b+':'+a)`   78gtyf63we4513
                            # print(c,end=' ')
                            # print('\n')

            output.close()